# _*_ coding: utf-8 _*_
# @Time : 2021/10/17 20:19
# @Author : Mr.C
# @File : TF-IDF算法
# @Project : ML_algorithm

# 引入依赖
import numpy as np
import pandas as pd

# 定义数据和预处理
docA = "The cat sat on my bed"
docB = "The dog sat on my knees"

bowA = docA.split(" ")
bowB = docB.split(" ")
# print(bowA, bowB)

# 构建一个完整的词库
wordSet = set(bowA).union(set(bowB))
# print(wordSet)

# 进行词数统计
# 用统计字典保存词出现的次数
wordDictA = dict.fromkeys(wordSet, 0)
wordDictB = dict.fromkeys(wordSet, 0)
# print(wordDictA, wordDictB)

# 遍历文档统计次数
for word in bowA:
    wordDictA[word] += 1

for word in bowB:
    wordDictB[word] += 1

# print(wordDictA, wordDictB)
# df = pd.DataFrame([wordDictA, wordDictB])
# print(df)

# 计算TF词频
def computeTF(wordDict, bow):
    # 用一个字典对象记录所有的TF,把所有的词在bow文档里的tf都计算出来
    tfDict = {}
    nbowCount = len(bow)

    for word ,count in wordDict.items():
        tfDict[word] = count / nbowCount
    return tfDict

tfA = computeTF(wordDictA, bowA)
tfB = computeTF(wordDictB, bowB)
# print(tfA)

# 计算逆文档频率IDF
def computeIDF(wordDictList):
    # 用一个字典对象保存idf结果，每个词作为key,初始值为0
    idfDict = dict.fromkeys(wordDictList[0], 0)
    N = len(wordDictList)
    import math

    for wordDict in wordDictList:
        # 遍历字典中的每个词汇,统计Ni
        for word, count in wordDict.items():
            if count > 0:
                # 先把Ni增加1存入到idfDict
                idfDict[word] += 1

    # 已经得到所有的词汇i对应的Ni，现在根据公式把它替换成为idf值
    for word, ni in idfDict.items():
        idfDict[word] = math.log10((N + 1) / (ni + 1))

    return idfDict

idfs = computeIDF([wordDictA, wordDictB])
# print(idfs)

# 计算TF-IDf
def computeTFIDF(tf, idfs):
    tfidf = {}
    for word, tfval in tf.items():
        tfidf[word] = tfval * idfs[word]
    return tfidf

tfidfA = computeTFIDF(tfA, idfs)
tfidfB = computeTFIDF(tfB, idfs)

df = pd.DataFrame([tfidfA, tfidfB])
print(df)